
from qdrant_client import QdrantClient
import pdfplumber
from qdrant_client.http.models import VectorParams, Distance, Batch

from EmbeddingHelper import EmbeddingHelper
from PDFHelper import PDFHelper


class QDrantHelper:
    """
    QDrantHelper
    """
    HOSTNAME = "192.168.60.131"
    PORT = 6333
    QDRANT_SIZE = 1536
    qdrantClient = None
    QDRANT_SEARCH_COUNT = 3

    def __init__(self):
        # 创建连接
        self.qdrantClient = QdrantClient(url=f"http://{self.HOSTNAME}:{self.PORT}",timeout=10000)
        self.pdfHelper = PDFHelper()
        self.vectorHelper = EmbeddingHelper()

    def addVector(self,collectionName,payloads,vectors):
        self.qdrantClient.upsert(collection_name=collectionName,wait=True,points=
            Batch(ids=[i for i in range(len(vectors))],payloads=payloads,vectors=vectors))
        return True


    def file_to_vector(self,file_path,file_name):

        try:
            collection_info = self.qdrantClient.get_collection(collection_name=file_name)
        except Exception as e:
            collection_info = None
            print(e)

        #判断是否已经存储过此文件的向量
        if collection_info == None:
            collection_info = self.qdrantClient.create_collection(collection_name=file_name, vectors_config=VectorParams(size = self.QDRANT_SIZE,distance = Distance.COSINE))

            # pdf_content = self.pdfHelper.extract_all_text(file_path)
            docs = self.pdfHelper.pdf_to_docs(file_path, file_name)
            docs = self.pdfHelper.docs_to_chunks(docs)

            texts = [doc.page_content for doc in docs]
            metadata = [doc.metadata for doc in docs]

            payloads = [{"page_content": text, "metadata": metadata} for text, metadata in zip(texts, metadata)]

            embeddings = self.vectorHelper.getVector(texts)

            if self.addVector(file_name, payloads, embeddings):
                return file_path
        else:
            return file_path

    def query_retrival(self,user_input,collection_names,maxCount,chat_histories):
        # 获取用户输入的向量
        user_input_embedding = self.vectorHelper.getVector(user_input)

        # 查询
        similar_points = []
        for collection_name in collection_names:
            collection_info = self.qdrantClient.get_collection(collection_name=collection_name)
            if collection_info != None:
                scored_points_by_current_collection = self.qdrantClient.search(collection_name=collection_name,query_vector=user_input_embedding[0],limit=self.QDRANT_SEARCH_COUNT,with_payload=True)
                if len(scored_points_by_current_collection) > 0:
                    similar_points.extend(scored_points_by_current_collection)

        # 排序
        points = []
        for similar_point in similar_points:
            point = {
                "id":similar_point.id,
                "payload":similar_point.payload,
                "score":similar_point.score
            }
            points.append(point)

        points.sort(key=lambda x:x["score"],reverse=True)
        points = points[:maxCount]

        #构建背景信息
        context = ""
        for index,point in enumerate(points):
            context += f'背景信息 {index} , {point["payload"]["page_content"]}\n\n'

        chat_history_str = ""
        #构建历史对话设置
        for chat_item in chat_histories:
            if chat_item[0] != None:
                chat_history_str += f'user: {chat_item[0]}\n'
            if chat_item[1] != None:
                chat_history_str += f'assistant: {chat_item[1]}\n'

        prompt = f"""你善于根据文档内容和历史对话进行分析总结，可以基于`文档内容`和`对话历史`回答user的问题。但请注意：如果user提出的问题与`文档内容`无关，你可以不参考`文档内容`直接进行回答。
        文档内容：```
        {context}```
        
        对话历史：```
        {chat_history_str}```
        
        user:``` {user_input}```
        assistant:
        """

        return prompt